import paho.mqtt.client as mqtt
import time
import threading
import random
import logging
import json
import os
from typing import Dict, List, Callable, Optional, Tuple, Any
from dataclasses import dataclass
from datetime import datetime

# 配置日志系统
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("mqtt_test_bed.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("MQTTTestBed")


@dataclass
class TestResult:
    """测试结果数据类，用于标准化测试输出"""
    test_name: str
    start_time: datetime
    end_time: datetime
    success: bool
    message: str
    metrics: Dict[str, Any] = None
    errors: List[str] = None

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典，便于序列化"""
        return {
            "test_name": self.test_name,
            "start_time": self.start_time.isoformat(),
            "end_time": self.end_time.isoformat(),
            "duration": (self.end_time - self.start_time).total_seconds(),
            "success": self.success,
            "message": self.message,
            "metrics": self.metrics or {},
            "errors": self.errors or []
        }


class MQTTTestBed:
    def __init__(self, broker_host: str = "localhost", broker_port: int = 1883, 
                 keepalive: int = 60, config_file: str = None):
        """
        初始化MQTT测试床
        :param broker_host: MQTT broker主机地址
        :param broker_port: MQTT broker端口
        :param keepalive: 保持连接时间(秒)
        :param config_file: 配置文件路径(可选)
        """
        # 从配置文件加载配置（如果提供）
        self.config = self._load_config(config_file) if config_file else {}
        
        # 配置优先级：参数 > 配置文件 > 默认值
        self.broker_host = broker_host or self.config.get('broker_host', "localhost")
        self.broker_port = broker_port or self.config.get('broker_port', 1883)
        self.keepalive = keepalive or self.config.get('keepalive', 60)
        
        # 存储订阅者收到的消息
        self.received_messages: Dict[str, List[Tuple[str, float]]] = {}  # (payload, timestamp)
        
        # 客户端列表
        self.clients: Dict[str, mqtt.Client] = {}
        
        # 连接状态
        self.connection_status: Dict[str, bool] = {}
        
        # 测试结果记录
        self.test_results: List[TestResult] = []
        
        # 事件回调
        self.event_callbacks: Dict[str, List[Callable]] = {
            'on_message': [],
            'on_connect': [],
            'on_disconnect': []
        }
        
        logger.info(f"MQTT测试床初始化完成，连接到 broker: {self.broker_host}:{self.broker_port}")

    def _load_config(self, config_file: str) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(config_file, 'r') as f:
                return json.load(f)
        except Exception as e:
            logger.warning(f"加载配置文件失败: {str(e)}，将使用默认配置")
            return {}

    def save_config(self, config_file: str) -> bool:
        """保存当前配置到文件"""
        try:
            config = {
                'broker_host': self.broker_host,
                'broker_port': self.broker_port,
                'keepalive': self.keepalive
            }
            with open(config_file, 'w') as f:
                json.dump(config, f, indent=2)
            logger.info(f"配置已保存到 {config_file}")
            return True
        except Exception as e:
            logger.error(f"保存配置文件失败: {str(e)}")
            return False

    def register_callback(self, event: str, callback: Callable):
        """
        注册事件回调函数
        :param event: 事件类型 ('on_message', 'on_connect', 'on_disconnect')
        :param callback: 回调函数
        """
        if event in self.event_callbacks:
            self.event_callbacks[event].append(callback)
            logger.info(f"已注册 {event} 事件回调")
        else:
            logger.warning(f"未知事件类型: {event}")

    def on_connect(self, client_id: str, client: mqtt.Client, userdata, flags, rc):
        """连接回调函数"""
        if rc == 0:
            self.connection_status[client_id] = True
            logger.info(f"客户端 {client_id} 连接成功")
        else:
            self.connection_status[client_id] = False
            logger.error(f"客户端 {client_id} 连接失败，错误代码: {rc}")
        
        # 触发回调
        for callback in self.event_callbacks['on_connect']:
            try:
                callback(client_id, rc == 0, rc)
            except Exception as e:
                logger.error(f"连接回调执行失败: {str(e)}")

    def on_disconnect(self, client_id: str, client: mqtt.Client, userdata, rc):
        """断开连接回调函数"""
        self.connection_status[client_id] = False
        logger.info(f"客户端 {client_id} 断开连接，错误代码: {rc}")
        
        # 触发回调
        for callback in self.event_callbacks['on_disconnect']:
            try:
                callback(client_id, rc)
            except Exception as e:
                logger.error(f"断开连接回调执行失败: {str(e)}")

    def on_message(self, client_id: str, client: mqtt.Client, msg):
        """消息接收回调函数"""
        topic = msg.topic
        payload = msg.payload.decode()
        timestamp = time.time()
        
        if topic not in self.received_messages:
            self.received_messages[topic] = []
        
        self.received_messages[topic].append((payload, timestamp))
        logger.debug(f"客户端 {client_id} 收到消息 - 主题: {topic}, 内容: {payload}")
        
        # 触发回调
        for callback in self.event_callbacks['on_message']:
            try:
                callback(client_id, topic, payload, timestamp)
            except Exception as e:
                logger.error(f"消息回调执行失败: {str(e)}")

    def create_client(self, client_id: str, clean_session: bool = True, 
                      username: Optional[str] = None, password: Optional[str] = None,
                      will_topic: Optional[str] = None, will_payload: Optional[str] = None,
                      will_qos: int = 0, will_retain: bool = False) -> mqtt.Client:
        """
        创建一个新的MQTT客户端，支持设置遗嘱消息
        :param client_id: 客户端ID
        :param clean_session: 是否使用干净会话
        :param username: 用户名(可选)
        :param password: 密码(可选)
        :param will_topic: 遗嘱消息主题(可选)
        :param will_payload: 遗嘱消息内容(可选)
        :param will_qos: 遗嘱消息QoS
        :param will_retain: 遗嘱消息是否保留
        """
        if client_id in self.clients:
            logger.warning(f"客户端 {client_id} 已存在，将被替换")
        
        # 创建客户端
        client = mqtt.Client(client_id=client_id, clean_session=clean_session)
        
        # 设置用户名和密码
        if username and password:
            client.username_pw_set(username, password)
        
        # 设置遗嘱消息
        if will_topic and will_payload:
            client.will_set(will_topic, will_payload, will_qos, will_retain)
            logger.info(f"客户端 {client_id} 已设置遗嘱消息: {will_topic}")
        
        # 绑定回调函数，使用lambda传递client_id
        client.on_connect = lambda client, userdata, flags, rc: self.on_connect(client_id, client, userdata, flags, rc)
        client.on_disconnect = lambda client, userdata, rc: self.on_disconnect(client_id, client, userdata, rc)
        client.on_message = lambda client, userdata, msg: self.on_message(client_id, client, msg)
        
        # 存储客户端
        self.clients[client_id] = client
        self.connection_status[client_id] = False
        
        logger.info(f"客户端 {client_id} 创建成功")
        return client

    def connect_client(self, client_id: str, timeout: int = 5) -> bool:
        """
        连接客户端到broker
        :param client_id: 客户端ID
        :param timeout: 连接超时时间(秒)
        """
        if client_id not in self.clients:
            logger.error(f"客户端 {client_id} 不存在")
            return False
        
        client = self.clients[client_id]
        start_time = time.time()
        
        try:
            client.connect(self.broker_host, self.broker_port, self.keepalive)
            # 启动网络循环线程
            client.loop_start()
            
            # 等待连接完成
            while not self.connection_status.get(client_id, False) and (time.time() - start_time) < timeout:
                time.sleep(0.1)
                
            return self.connection_status[client_id]
        except Exception as e:
            logger.error(f"客户端 {client_id} 连接异常: {str(e)}")
            return False

    def disconnect_client(self, client_id: str):
        """
        断开客户端连接
        :param client_id: 客户端ID
        """
        if client_id not in self.clients:
            logger.error(f"客户端 {client_id} 不存在")
            return
        
        client = self.clients[client_id]
        try:
            client.disconnect()
            client.loop_stop()
            logger.info(f"客户端 {client_id} 已断开连接")
        except Exception as e:
            logger.error(f"客户端 {client_id} 断开连接异常: {str(e)}")

    def subscribe(self, client_id: str, topic: str, qos: int = 0) -> bool:
        """
        订阅主题
        :param client_id: 客户端ID
        :param topic: 要订阅的主题
        :param qos: 服务质量等级
        """
        if client_id not in self.clients:
            logger.error(f"客户端 {client_id} 不存在")
            return False
        
        if not self.connection_status[client_id]:
            logger.error(f"客户端 {client_id} 未连接，无法订阅")
            return False
        
        client = self.clients[client_id]
        try:
            result, mid = client.subscribe(topic, qos)
            success = result == 0
            if success:
                logger.info(f"客户端 {client_id} 订阅主题 {topic} 成功")
            else:
                logger.error(f"客户端 {client_id} 订阅主题 {topic} 失败")
            return success
        except Exception as e:
            logger.error(f"客户端 {client_id} 订阅异常: {str(e)}")
            return False

    def unsubscribe(self, client_id: str, topic: str) -> bool:
        """
        取消订阅主题
        :param client_id: 客户端ID
        :param topic: 要取消订阅的主题
        """
        if client_id not in self.clients:
            logger.error(f"客户端 {client_id} 不存在")
            return False
        
        if not self.connection_status[client_id]:
            logger.error(f"客户端 {client_id} 未连接，无法取消订阅")
            return False
        
        client = self.clients[client_id]
        try:
            result, mid = client.unsubscribe(topic)
            success = result == 0
            if success:
                logger.info(f"客户端 {client_id} 取消订阅主题 {topic} 成功")
            else:
                logger.error(f"客户端 {client_id} 取消订阅主题 {topic} 失败")
            return success
        except Exception as e:
            logger.error(f"客户端 {client_id} 取消订阅异常: {str(e)}")
            return False

    def publish(self, client_id: str, topic: str, payload: str, qos: int = 0, 
                retain: bool = False, timeout: int = 10) -> bool:
        """
        发布消息
        :param client_id: 客户端ID
        :param topic: 发布的主题
        :param payload: 消息内容
        :param qos: 服务质量等级
        :param retain: 是否保留消息
        :param timeout: 发布超时时间(秒)
        """
        if client_id not in self.clients:
            logger.error(f"客户端 {client_id} 不存在")
            return False
        
        if not self.connection_status[client_id]:
            logger.error(f"客户端 {client_id} 未连接，无法发布")
            return False
        
        client = self.clients[client_id]
        try:
            result = client.publish(topic, payload, qos, retain)
            # 等待消息发布完成
            result.wait_for_publish(timeout)
            success = result.rc == 0
            
            if success:
                logger.debug(f"客户端 {client_id} 发布消息到 {topic} 成功")
            else:
                logger.error(f"客户端 {client_id} 发布消息到 {topic} 失败，错误代码: {result.rc}")
            return success
        except Exception as e:
            logger.error(f"客户端 {client_id} 发布异常: {str(e)}")
            return False

    def get_received_messages(self, topic: Optional[str] = None, 
                             since: Optional[float] = None) -> Dict[str, List[Tuple[str, float]]]:
        """
        获取收到的消息
        :param topic: 可选，指定主题，None则返回所有主题的消息
        :param since: 可选，指定时间戳，只返回该时间之后的消息
        """
        result = {}
        
        if topic:
            if topic in self.received_messages:
                messages = self.received_messages[topic]
                if since:
                    messages = [(p, t) for p, t in messages if t >= since]
                result[topic] = messages
        else:
            for t, messages in self.received_messages.items():
                if since:
                    filtered = [(p, ts) for p, ts in messages if ts >= since]
                    if filtered:
                        result[t] = filtered
                else:
                    result[t] = messages
        
        return result.copy()

    def clear_received_messages(self):
        """清除已接收的消息"""
        self.received_messages.clear()
        logger.info("已清除所有接收的消息")

    def save_test_results(self, file_path: str = "test_results.json") -> bool:
        """保存测试结果到JSON文件"""
        try:
            results = [r.to_dict() for r in self.test_results]
            with open(file_path, 'w') as f:
                json.dump(results, f, indent=2)
            logger.info(f"测试结果已保存到 {file_path}")
            return True
        except Exception as e:
            logger.error(f"保存测试结果失败: {str(e)}")
            return False

    def run_basic_test(self, test_name: str = "basic_mqtt_test", 
                      test_topic: str = "test/mqtt/testbed",
                      qos_level: int = 1) -> TestResult:
        """
        运行基本的MQTT功能测试
        :param test_name: 测试名称
        :param test_topic: 测试使用的主题
        :param qos_level: 测试使用的QoS级别
        :return: 测试结果对象
        """
        start_time = datetime.now()
        errors = []
        metrics = {}
        success = False
        message = ""
        
        logger.info(f"\n===== 开始基本MQTT功能测试: {test_name} =====")
        
        try:
            # 1. 创建测试客户端
            pub_client_id = f"{test_name}_publisher"
            sub_client_id = f"{test_name}_subscriber"
            
            self.create_client(pub_client_id)
            self.create_client(sub_client_id)
            
            # 2. 连接客户端
            pub_connected = self.connect_client(pub_client_id)
            sub_connected = self.connect_client(sub_client_id)
            
            if not pub_connected:
                raise Exception(f"发布者客户端 {pub_client_id} 连接失败")
            if not sub_connected:
                raise Exception(f"订阅者客户端 {sub_client_id} 连接失败")
            
            # 3. 订阅测试主题
            sub_success = self.subscribe(sub_client_id, test_topic, qos_level)
            
            if not sub_success:
                raise Exception(f"订阅主题 {test_topic} 失败")
            
            # 等待订阅生效
            time.sleep(1)
            
            # 4. 发布测试消息
            test_messages = [
                "测试消息 1",
                "测试消息 2",
                "这是一条带特殊字符的消息: !@#$%^&*()",
                str(random.randint(1000, 9999))  # 随机数字消息
            ]
            
            publish_start = time.time()
            publish_success_count = 0
            
            for msg in test_messages:
                if self.publish(pub_client_id, test_topic, msg, qos_level):
                    publish_success_count += 1
                time.sleep(0.5)  # 等待消息传递
            
            publish_time = time.time() - publish_start
            metrics["publish_time"] = publish_time
            metrics["publish_success_rate"] = publish_success_count / len(test_messages)
            
            # 等待消息接收
            time.sleep(2)
            
            # 5. 验证接收的消息
            received = self.get_received_messages(test_topic)
            received_messages = received.get(test_topic, [])
            received_payloads = [p for p, _ in received_messages]
            
            metrics["sent_messages"] = len(test_messages)
            metrics["received_messages"] = len(received_messages)
            
            logger.info(f"\n验证结果:")
            logger.info(f"发送的消息数: {len(test_messages)}")
            logger.info(f"接收的消息数: {len(received_messages)}")
            
            if len(test_messages) == len(received_messages):
                match = all(msg in received_payloads for msg in test_messages)
                if match:
                    success = True
                    message = "测试成功: 所有消息都被正确接收"
                else:
                    errors.append("消息内容不匹配")
                    message = "测试失败: 消息内容不匹配"
            else:
                errors.append(f"消息数量不匹配 (发送: {len(test_messages)}, 接收: {len(received_messages)})")
                message = "测试失败: 消息数量不匹配"
                
        except Exception as e:
            errors.append(str(e))
            message = f"测试执行过程中发生错误: {str(e)}"
            logger.error(message, exc_info=True)
            
        finally:
            # 6. 清理资源
            if 'pub_client_id' in locals():
                self.disconnect_client(pub_client_id)
            if 'sub_client_id' in locals():
                self.disconnect_client(sub_client_id)
            self.clear_received_messages()
        
        end_time = datetime.now()
        test_result = TestResult(
            test_name=test_name,
            start_time=start_time,
            end_time=end_time,
            success=success,
            message=message,
            metrics=metrics,
            errors=errors
        )
        
        self.test_results.append(test_result)
        logger.info(f"===== 基本MQTT功能测试结束: {test_name} =====")
        
        return test_result

    def run_stress_test(self, num_messages: int = 100, message_size: int = 1024,
                       test_name: str = "mqtt_stress_test",
                       topic: str = "test/mqtt/stresstest",
                       qos_level: int = 0) -> TestResult:
        """
        运行MQTT压力测试
        :param num_messages: 消息数量
        :param message_size: 每条消息的大小(字节)
        :param test_name: 测试名称
        :param topic: 测试使用的主题
        :param qos_level: 测试使用的QoS级别
        :return: 测试结果对象
        """
        start_time = datetime.now()
        errors = []
        metrics = {}
        success = False
        message = ""
        
        logger.info(f"\n===== 开始MQTT压力测试: {test_name} ({num_messages}条消息，每条{message_size}字节) =====")
        
        try:
            # 创建客户端
            pub_client_id = f"{test_name}_publisher"
            sub_client_id = f"{test_name}_subscriber"
            
            self.create_client(pub_client_id)
            self.create_client(sub_client_id)
            
            # 连接客户端
            if not (self.connect_client(pub_client_id) and self.connect_client(sub_client_id)):
                raise Exception("客户端连接失败")
            
            # 订阅主题
            if not self.subscribe(sub_client_id, topic, qos_level):
                raise Exception(f"订阅主题 {topic} 失败")
            
            time.sleep(1)
            
            # 生成测试消息
            test_message_content = "X" * message_size  # 生成固定大小的消息
            
            # 记录开始时间
            test_start_time = time.time()
            
            # 发送消息
            success_count = 0
            for i in range(num_messages):
                payload = f"{i}:{test_message_content}"
                if self.publish(pub_client_id, topic, payload, qos_level):
                    success_count += 1
                
                # 每发送10%的消息打印一次进度
                if (i + 1) % max(1, num_messages // 10) == 0:
                    logger.info(f"已发送 {i + 1}/{num_messages} 条消息")
            
            # 等待所有消息接收
            time.sleep(5)
            
            # 计算耗时
            elapsed_time = time.time() - test_start_time
            
            # 获取接收的消息数
            received = self.get_received_messages(topic)
            received_count = len(received.get(topic, []))
            
            # 计算指标
            metrics = {
                "total_messages": num_messages,
                "successfully_sent": success_count,
                "successfully_received": received_count,
                "elapsed_time_seconds": elapsed_time,
                "send_rate_messages_per_sec": success_count / elapsed_time if elapsed_time > 0 else 0,
                "receive_rate_messages_per_sec": received_count / elapsed_time if elapsed_time > 0 else 0,
                "message_size_bytes": message_size,
                "total_data_kb": num_messages * message_size / 1024,
                "loss_rate": 1 - (received_count / success_count) if success_count > 0 else 1
            }
            
            # 判断测试是否成功（可以根据需要调整成功标准）
            success = received_count >= success_count * 0.9  # 允许10%的消息丢失
            message = "压力测试成功" if success else "压力测试失败: 消息丢失率过高"
            
            # 打印结果
            logger.info("\n压力测试结果:")
            for key, value in metrics.items():
                logger.info(f"{key}: {value:.2f}" if isinstance(value, float) else f"{key}: {value}")
                
        except Exception as e:
            errors.append(str(e))
            message = f"压力测试执行过程中发生错误: {str(e)}"
            logger.error(message, exc_info=True)
            
        finally:
            # 清理资源
            if 'pub_client_id' in locals():
                self.disconnect_client(pub_client_id)
            if 'sub_client_id' in locals():
                self.disconnect_client(sub_client_id)
            self.clear_received_messages()
        
        end_time = datetime.now()
        test_result = TestResult(
            test_name=test_name,
            start_time=start_time,
            end_time=end_time,
            success=success,
            message=message,
            metrics=metrics,
            errors=errors
        )
        
        self.test_results.append(test_result)
        logger.info(f"===== MQTT压力测试结束: {test_name} =====")
        
        return test_result

    def run_will_message_test(self, test_name: str = "will_message_test",
                             will_topic: str = "test/mqtt/will") -> TestResult:
        """
        测试遗嘱消息功能
        :param test_name: 测试名称
        :param will_topic: 遗嘱消息主题
        :return: 测试结果对象
        """
        start_time = datetime.now()
        errors = []
        metrics = {}
        success = False
        message = ""
        
        logger.info(f"\n===== 开始遗嘱消息测试: {test_name} =====")
        
        try:
            # 创建订阅者客户端（订阅遗嘱主题）
            sub_client_id = f"{test_name}_subscriber"
            self.create_client(sub_client_id)
            if not self.connect_client(sub_client_id):
                raise Exception(f"订阅者客户端 {sub_client_id} 连接失败")
            
            # 订阅遗嘱主题
            if not self.subscribe(sub_client_id, will_topic):
                raise Exception(f"订阅遗嘱主题 {will_topic} 失败")
            
            # 创建带有遗嘱的客户端
            will_client_id = f"{test_name}_will_client"
            will_payload = f"客户端 {will_client_id} 异常断开连接"
            self.create_client(
                will_client_id,
                will_topic=will_topic,
                will_payload=will_payload,
                will_qos=1
            )
            
            if not self.connect_client(will_client_id):
                raise Exception(f"遗嘱客户端 {will_client_id} 连接失败")
            
            # 等待连接稳定
            time.sleep(2)
            
            # 异常断开连接（模拟客户端崩溃）
            client = self.clients[will_client_id]
            client.socket().close()  # 强制关闭socket，不发送DISCONNECT包
            logger.info(f"已强制断开 {will_client_id} 的连接，等待遗嘱消息")
            
            # 等待遗嘱消息
            time.sleep(5)
            
            # 检查是否收到遗嘱消息
            received = self.get_received_messages(will_topic)
            received_messages = [p for p, _ in received.get(will_topic, [])]
            
            metrics["expected_will_payload"] = will_payload
            metrics["received_will_count"] = len(received_messages)
            
            if will_payload in received_messages:
                success = True
                message = "遗嘱消息测试成功: 收到预期的遗嘱消息"
            else:
                errors.append("未收到预期的遗嘱消息")
                message = "遗嘱消息测试失败: 未收到预期的遗嘱消息"
                
        except Exception as e:
            errors.append(str(e))
            message = f"遗嘱消息测试执行过程中发生错误: {str(e)}"
            logger.error(message, exc_info=True)
            
        finally:
            # 清理资源
            if 'sub_client_id' in locals():
                self.disconnect_client(sub_client_id)
            if 'will_client_id' in locals() and will_client_id in self.clients:
                self.disconnect_client(will_client_id)
            self.clear_received_messages()
        
        end_time = datetime.now()
        test_result = TestResult(
            test_name=test_name,
            start_time=start_time,
            end_time=end_time,
            success=success,
            message=message,
            metrics=metrics,
            errors=errors
        )
        
        self.test_results.append(test_result)
        logger.info(f"===== 遗嘱消息测试结束: {test_name} =====")
        
        return test_result

    def run_all_tests(self, output_file: str = "all_test_results.json") -> List[TestResult]:
        """运行所有测试用例"""
        logger.info("\n===== 开始运行所有测试用例 =====")
        
        results = []
        results.append(self.run_basic_test())
        results.append(self.run_stress_test())
        results.append(self.run_will_message_test())
        
        self.save_test_results(output_file)
        logger.info(f"===== 所有测试用例运行结束 =====")
        
        return results


# 示例用法
if __name__ == "__main__":
    # 初始化测试床，默认连接到本地MQTT broker
    # 如果你的broker在其他位置，可以直接指定或使用配置文件
    test_bed = MQTTTestBed(broker_host="localhost", broker_port=1883)
    
    # 可选：保存当前配置
    test_bed.save_config("mqtt_test_config.json")
    
    # 运行所有测试
    test_bed.run_all_tests()
    
    # 或者单独运行特定测试
    # test_bed.run_basic_test()
    # test_bed.run_stress_test(num_messages=500, message_size=2048)
    # test_bed.run_will_message_test()
    
    # 查看测试结果
    logger.info("\n测试总结:")
    for result in test_bed.test_results:
        status = "成功" if result.success else "失败"
        logger.info(f"{result.test_name}: {status} ({result.message})")
